import torch
import torch.nn as nn
import torch.nn.functional as F


class CIFAR10CNN(nn.Module):
    """基于CNN的CIFAR-10图像分类模型"""
    def __init__(self, num_classes=10):
        super().__init__()
        
        # 卷积块1：输入3通道，输出32通道
        self.conv_block1 = nn.Sequential(
            # 卷积层1：3→32，3×3卷积， padding=1保持尺寸
            nn.Conv2d(in_channels=3, out_channels=32, kernel_size=3, padding=1),
            nn.BatchNorm2d(32),  # 批归一化
            nn.ReLU(inplace=True),  # ReLU激活
            
            # 卷积层2：32→32，3×3卷积
            nn.Conv2d(in_channels=32, out_channels=32, kernel_size=3, padding=1),
            nn.BatchNorm2d(32),
            nn.ReLU(inplace=True),
            
            # 最大池化：2×2，步长2，输出尺寸减半
            nn.MaxPool2d(kernel_size=2, stride=2)
        )
        
        # 卷积块2：输入32通道，输出64通道
        self.conv_block2 = nn.Sequential(
            nn.Conv2d(32, 64, kernel_size=3, padding=1),
            nn.BatchNorm2d(64),
            nn.ReLU(inplace=True),
            
            nn.Conv2d(64, 64, kernel_size=3, padding=1),
            nn.BatchNorm2d(64),
            nn.ReLU(inplace=True),
            
            nn.MaxPool2d(2, 2)
        )
        
        # 卷积块3：输入64通道，输出128通道
        self.conv_block3 = nn.Sequential(
            nn.Conv2d(64, 128, kernel_size=3, padding=1),
            nn.BatchNorm2d(128),
            nn.ReLU(inplace=True),
            
            nn.Conv2d(128, 128, kernel_size=3, padding=1),
            nn.BatchNorm2d(128),
            nn.ReLU(inplace=True),
            
            nn.MaxPool2d(2, 2)
        )
        
        # 分类头：全连接层
        self.classifier = nn.Sequential(
            # 先将特征展平：经过3次池化后，32×32→4×4（32/(2^3)=4）
            # 展平尺寸：128 channels × 4×4 = 2048
            nn.Linear(128 * 4 * 4, 512),
            nn.ReLU(inplace=True),
            nn.Dropout(0.5),  # 防止过拟合
            
            nn.Linear(512, num_classes)
        )
        
        # 初始化权重
        self._initialize_weights()

    def forward(self, x):
        """前向传播过程"""
        # 特征提取
        x = self.conv_block1(x)  # [batch, 32, 16, 16]（32/2=16）
        x = self.conv_block2(x)  # [batch, 64, 8, 8]（16/2=8）
        x = self.conv_block3(x)  # [batch, 128, 4, 4]（8/2=4）
        
        # 展平特征图（保留batch维度）
        x = x.view(x.size(0), -1)  # [batch, 128*4*4=2048]
        
        # 分类输出
        x = self.classifier(x)  # [batch, num_classes]
        return x

    def _initialize_weights(self):
        """初始化卷积层和全连接层的权重"""
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                # 卷积层：Kaiming正态分布初始化（适用于ReLU激活）
                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
                if m.bias is not None:
                    nn.init.constant_(m.bias, 0)  # 偏置初始化为0
            elif isinstance(m, nn.BatchNorm2d):
                # 批归一化：权重初始化为1，偏置为0
                nn.init.constant_(m.weight, 1)
                nn.init.constant_(m.bias, 0)
            elif isinstance(m, nn.Linear):
                # 全连接层：正态分布初始化
                nn.init.normal_(m.weight, 0, 0.01)
                nn.init.constant_(m.bias, 0)


# 测试代码
if __name__ == "__main__":
    # 超参数设置
    batch_size = 16
    num_classes = 10
    img_size = (3, 32, 32)  # CIFAR-10图像尺寸：3通道，32×32
    
    # 初始化模型
    model = CIFAR10CNN(num_classes=num_classes)
    
    # 生成随机输入（模拟一批图像）
    x = torch.randn(batch_size, *img_size)  # [batch, channels, height, width]
    
    # 前向传播
    output = model(x)
    
    # 验证输出形状
    print(f"输入形状: {x.shape}")
    print(f"输出形状: {output.shape}")  # 应输出 [16, 10]
    
    # 计算模型参数量
    total_params = sum(p.numel() for p in model.parameters())
    print(f"模型总参数量: {total_params:,}")